import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;

public class AIAssistant {
    private static final String API_URL = "https://api.deepseek.com/v1/chat/completions";
    private String apiKey;

    public AIAssistant() {
        // 在这里填入您的DeepSeek API Key
        this.apiKey = "sk-f8eb2ff921ec4619a0d166e7b68ad269";

        // 如果没有填入API Key，尝试从环境变量读取
        if (this.apiKey.equals("sk-your-deepseek-api-key-here")) {
            String envKey = System.getenv("DEEPSEEK_API_KEY");
            if (envKey != null && !envKey.trim().isEmpty()) {
                this.apiKey = envKey;
            }
        }
    }

    public void setApiKey(String apiKey) {
        this.apiKey = apiKey;
    }

    public boolean isConfigured() {
        return apiKey != null && !apiKey.trim().isEmpty();
    }

    public String generateContent(String prompt, String context) {
        if (!isConfigured()) {
            return "请先配置DeepSeek API Key";
        }

        try {
            String fullPrompt = buildPrompt(prompt, context);
            String response = callDeepSeekAPI(fullPrompt);
            return parseResponse(response);
        } catch (Exception e) {
            e.printStackTrace();
            return "AI助手调用失败: " + e.getMessage();
        }
    }

    private String buildPrompt(String prompt, String context) {
        StringBuilder sb = new StringBuilder();

        if (context != null && !context.trim().isEmpty()) {
            sb.append("基于以下文本内容：\n");
            sb.append(context);
            sb.append("\n\n");
        }

        sb.append("请帮我：");
        sb.append(prompt);
        sb.append("\n\n请用中文回答，内容要有用且符合上下文。");

        return sb.toString();
    }

    private String callDeepSeekAPI(String prompt) throws Exception {
        URL url = new URL(API_URL);
        HttpURLConnection conn = (HttpURLConnection) url.openConnection();

        // 设置请求方法和头部
        conn.setRequestMethod("POST");
        conn.setRequestProperty("Content-Type", "application/json");
        conn.setRequestProperty("Authorization", "Bearer " + apiKey);
        conn.setDoOutput(true);

        // 构建JSON请求体（手动拼接）
        String jsonRequest = buildJsonRequest(prompt);

        // 发送请求
        try (OutputStream os = conn.getOutputStream()) {
            byte[] input = jsonRequest.getBytes(StandardCharsets.UTF_8);
            os.write(input, 0, input.length);
        }

        // 读取响应
        int responseCode = conn.getResponseCode();
        InputStream inputStream = responseCode == 200 ? conn.getInputStream() : conn.getErrorStream();

        StringBuilder response = new StringBuilder();
        try (BufferedReader br = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {
            String line;
            while ((line = br.readLine()) != null) {
                response.append(line);
            }
        }

        if (responseCode != 200) {
            throw new RuntimeException("API调用失败，状态码: " + responseCode + ", 响应: " + response.toString());
        }

        return response.toString();
    }

    private String buildJsonRequest(String prompt) {
        // 手动构建JSON，确保字符串正确转义
        String escapedPrompt = escapeJson(prompt);
        return String.format(
                "{\"model\":\"deepseek-chat\",\"messages\":[{\"role\":\"user\",\"content\":\"%s\"}],\"max_tokens\":1000,\"temperature\":0.7}",
                escapedPrompt);
    }

    private String escapeJson(String str) {
        if (str == null)
            return "";
        return str.replace("\\", "\\\\")
                .replace("\"", "\\\"")
                .replace("\n", "\\n")
                .replace("\r", "\\r")
                .replace("\t", "\\t");
    }

    private String parseResponse(String response) {
        try {
            // 简单的JSON解析，查找content字段的内容
            String contentMarker = "\"content\":\"";
            int startIndex = response.indexOf(contentMarker);
            if (startIndex == -1) {
                return "AI助手没有返回有效内容";
            }

            startIndex += contentMarker.length();
            int endIndex = findJsonStringEnd(response, startIndex);

            if (endIndex == -1) {
                return "解析AI响应失败";
            }

            String content = response.substring(startIndex, endIndex);
            return unescapeJson(content);

        } catch (Exception e) {
            e.printStackTrace();
            return "解析AI响应失败: " + e.getMessage();
        }
    }

    private int findJsonStringEnd(String json, int startIndex) {
        boolean escaped = false;
        for (int i = startIndex; i < json.length(); i++) {
            char c = json.charAt(i);
            if (escaped) {
                escaped = false;
                continue;
            }
            if (c == '\\') {
                escaped = true;
            } else if (c == '"') {
                return i;
            }
        }
        return -1;
    }

    private String unescapeJson(String str) {
        if (str == null)
            return "";
        return str.replace("\\\"", "\"")
                .replace("\\\\", "\\")
                .replace("\\n", "\n")
                .replace("\\r", "\r")
                .replace("\\t", "\t");
    }

    // 预设的写作助手功能
    public String improveText(String text) {
        return generateContent("帮我改进这段文字，让它更加清晰和优雅", text);
    }

    public String summarizeText(String text) {
        return generateContent("请为这段文字写一个简要总结", text);
    }

    public String expandText(String text) {
        return generateContent("请扩展这段文字，添加更多细节和描述", text);
    }

    public String translateText(String text, String targetLanguage) {
        return generateContent("请将这段文字翻译成" + targetLanguage, text);
    }

    public String generateTitle(String text) {
        return generateContent("请为这段文字生成一个合适的标题", text);
    }

    public String checkGrammar(String text) {
        return generateContent("请检查这段文字的语法和表达，并提出改进建议", text);
    }
}